import numpy as np
import pandas as pd
import statsmodels.api as sm

#gain因子 参考广发v型处置效应研报
def gain(close,volume,period):
    if not (isinstance(close,np.ndarray) and isinstance(volume,np.ndarray)):
        raise TypeError('只接受np.array数据')
    elif len(close.shape)!=2 or len(volume.shape)!=2:
        raise ValueError('只接受n*1格式的array')
    elif close.shape[1]!=1 or volume.shape[1]!=1:
        raise ValueError('只接受n*1格式的array')
    elif close.shape[0]!=volume.shape[0]:
        raise ValueError('格式不匹配!')
    elif close.shape[0]<period:
        raise ValueError('数据期数小于等于计算期数')
    
    n=close.shape[0]-period
    Gain=np.repeat(np.nan,n).reshape(-1,1)
    for i in range(n):
        Close=close[i:i+period].copy()
        Volume=volume[i:i+period].copy()
        Price=close[i+period]
        Close_change=(Price-Close)/Price
        Close_change=(np.abs(Close_change)+Close_change)/2
        w=Volume*((((1-Volume)[::-1]).cumprod())[::-1]).reshape(-1,1)/(1-Volume)
        w_norm=w/sum(w)
        Gain[i,0]=(w_norm*Close_change)[-1,0]
    return Gain

#光大rsrs右偏修正标准分因子
def rsrs(high,low,period_beta,period_std):
    if not (isinstance(high,np.ndarray) and isinstance(low,np.ndarray)):
        raise TypeError('只接受np.array数据')
    elif len(high.shape)!=2 or len(low.shape)!=2:
        raise ValueError('只接受n*1格式的array')
    elif high.shape[1]!=1 or low.shape[1]!=1:
        raise ValueError('只接受n*1格式的array')
    elif high.shape[0]!=low.shape[0]:
        raise ValueError('格式不匹配!')
    elif high.shape[0]<period_beta+period_std:
        raise ValueError('数据期数小于等于计算期数')
    elif sum(pd.isnull(high)+pd.isnull(low))>0:
        raise ValueError('存在缺失数据!')
    
    n_beta=high.shape[0]-period_beta+1
    rsquare=np.repeat(np.nan,n_beta).reshape(-1,1)
    beta=np.repeat(np.nan,n_beta).reshape(-1,1)
    for i in range(n_beta):
        High=high[i:i+period_beta].copy()
        Low=low[i:i+period_beta].copy()
        x=sm.add_constant(Low)
        model=sm.OLS(High,x)
        result=model.fit()
        beta[i,0]=result.params[-1]
        rsquare[i,0]=result.rsquared
    
    beta_tmp=pd.Series(beta.reshape(-1))
    mean=beta_tmp.rolling(window=period_std).mean()
    std=beta_tmp.rolling(window=period_std).std()
    beta_std=(beta_tmp-mean)/std
    beta_std=beta_std.values.reshape(-1,1)

    score=beta*beta_std*rsquare
    return score


#测试
if __name__=='__main__':
    #Gain测试用例
    close=np.array([51,52.5,49,48,48.5]).reshape(-1,1)
    price_t=48.5
    volume=np.array([0.01,0.005,0.08,0.02,0.01]).reshape(-1,1)
    answer=0.00184
    a=gain(close,volume,4)
    if abs(a[0,0]-answer)<0.0001:
        print('gain测试成功')
    
    #rsrs测试用例
    high=np.array([1,2,1,2,3]).reshape(-1,1)
    low=np.array([1,2,3,4,5]).reshape(-1,1)
    period_beta=2
    period_std=3
    func_score=rsrs(high,low,period_beta,period_std)
    score=1/np.sqrt(3)
    if abs(func_score[-1]-score)<1e-4:
        print('rsrs通过测试')

